from ipex_llm.transformers import AutoModelForSpeechSeq2Seq
from transformers import WhisperProcessor
import librosa
import torch
import time

# Load model and processor
model = AutoModelForSpeechSeq2Seq.from_pretrained(pretrained_model_name_or_path="./whisper-tiny",
                                                  load_in_4bit=True)
processor = WhisperProcessor.from_pretrained(pretrained_model_name_or_path="./whisper-tiny")

# Load audio file
data_zh, sample_rate_zh = librosa.load("1.wav", sr=16000)

# Define Chinese transcribe task
forced_decoder_ids = processor.get_decoder_prompt_ids(language="chinese", task="transcribe")
model.config.forced_decoder_ids = forced_decoder_ids
# Generate input features
input_features = processor(data_zh, sampling_rate=sample_rate_zh, return_tensors="pt").input_features

# Generate predicted tokens
with torch.inference_mode():
    st = time.time()
    predicted_ids = model.generate(input_features)#, forced_decoder_ids=forced_decoder_ids)  # Pass forced_decoder_ids here
    end = time.time()
    print(f'Inference time: {end - st} s')

    # Decode predicted tokens to text
    transcribe_str = processor.batch_decode(predicted_ids, skip_special_tokens=True)

    print('-' * 20, 'Chinese Transcription', '-' * 20)
    print(transcribe_str[0])